package com.hzw.saas.web.shell.util;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Objects;

import com.jcraft.jsch.Session;
import cn.hutool.core.lang.SimpleCache;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.ssh.JschUtil;
import lombok.extern.slf4j.Slf4j;

/**
 * @author zzl
 * @since 06/07/2021
 */
@Slf4j
public enum JschEnhanceSessionPool {

    INSTANCE;

    /**
     * SSH会话池，key：host，value：Session对象
     */
    private final SimpleCache<String, Session> cache = new SimpleCache<>(new HashMap<>());

    public String getKey(String sshHost, int sshPort, String sshUser) {
        return StrUtil.format("{}@{}:{}", sshUser, sshHost, sshPort);
    }

    /**
     * 获取Session，不存在返回null
     *
     * @param key 键
     * @return Session
     */
    public Session get(String key) {
        return cache.get(key);
    }

    /**
     * 获得一个SSH跳板机会话，重用已经使用的会话
     *
     * @param sshHost 跳板机主机
     * @param sshPort 跳板机端口
     * @param sshUser 跳板机用户名
     * @param sshPass 跳板机密码
     * @return SSH会话
     */
    public Session getSession(String sshHost, int sshPort, String sshUser, String sshPass) {
        final String key = getKey(sshHost, sshPort, sshUser);
        return this.cache.get(key, () -> JschUtil.openSession(sshHost, sshPort, sshUser, sshPass));
    }

    /**
     * 获得一个SSH跳板机会话，重用已经使用的会话
     *
     * @param sshHost 跳板机主机
     * @param sshPort 跳板机端口
     * @param sshUser 跳板机用户名
     * @param sshPass 跳板机密码
     * @param sid     sessionId 需确保唯一
     * @return SSH会话
     */
    public Session getSession(String sshHost, int sshPort, String sshUser, String sshPass, String sid) {
        return this.cache.get(sid, () -> JschUtil.openSession(sshHost, sshPort, sshUser, sshPass));
    }


    /**
     * 获得一个SSH跳板机会话，重用已经使用的会话
     *
     * @param sshHost    跳板机主机
     * @param sshPort    跳板机端口
     * @param sshUser    跳板机用户名
     * @param prvkey     跳板机私钥路径
     * @param passphrase 跳板机私钥密码
     * @return SSH会话
     */
    public Session getSession(String sshHost, int sshPort, String sshUser, String prvkey, byte[] passphrase) {
        final String key = getKey(sshHost, sshPort, sshUser);
        return this.cache.get(key, () -> JschUtil.openSession(sshHost, sshPort, sshUser, prvkey, passphrase));
    }

    /**
     * 加入Session
     *
     * @param key     键
     * @param session Session
     */
    public void put(String key, Session session) {
        this.cache.put(key, session);
    }

    /**
     * 关闭SSH连接会话
     *
     * @param key   session的缓存key
     * @param retry 重试次数
     */
    public boolean close(String key, int retry) {
        Session session = get(key);
        retry = Math.max(retry, 1);
        for (; (retry > 0) && (Objects.nonNull(session) && session.isConnected()); retry--) {
            session.disconnect();
            ThreadUtil.sleep(JschEnhanceUtil.RETRY_INTERVAL);
        }
        if (Objects.isNull(session) || !session.isConnected()) {
            this.cache.remove(key);
            log.debug("关闭session连接完成, key: {}", key);
            return true;
        } else {
            log.debug("关闭session连接失败, key: {}", key);
            return false;
        }
    }

    /**
     * 移除指定Session
     *
     * @param session Session会话
     * @since 4.1.15
     */
    public void remove(Session session) {
        if (null != session) {
            final Iterator<Map.Entry<String, Session>> iterator = this.cache.iterator();
            Map.Entry<String, Session> entry;
            while (iterator.hasNext()) {
                entry = iterator.next();
                if (session.equals(entry.getValue())) {
                    iterator.remove();
                    break;
                }
            }
        }
    }

    /**
     * 关闭所有SSH连接会话
     */
    public void closeAll() {
        Session session;
        for (Map.Entry<String, Session> entry : this.cache) {
            session = entry.getValue();
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
        }
        cache.clear();
    }

}
